import os
import glob
import pypandoc
import markitdown
from fastapi import APIRouter, UploadFile, File, Form
from fastapi.responses import JSONResponse
from typing import List, Optional
import shutil
from pathlib import Path

router = APIRouter(
    prefix="/pandoc",
    tags=["文档转换"],
    responses={404: {"description": "Not found"}},
)

# 支持的输入格式
SUPPORTED_INPUT_FORMATS = [
    "md", "markdown", "docx", "doc", "odt", "rtf", "epub", "html", "htm",
    "tex", "latex", "org", "rst", "asciidoc", "adoc", "t2t", "textile",
    "wiki", "mediawiki", "twiki", "csv", "tsv", "json", "xml"
]

# 支持的输出格式
SUPPORTED_OUTPUT_FORMATS = [
    "md", "markdown", "docx", "odt", "rtf", "epub", "html", "htm",
    "tex", "latex", "org", "rst", "asciidoc", "adoc", "t2t", "textile",
    "wiki", "mediawiki", "twiki", "csv", "tsv", "json", "xml", "pdf"
]

# PDF特殊处理（Pandoc不能直接从PDF转换）
PDF_FORMAT = "pdf"

def get_file_extension(filename):
    """获取文件扩展名（不带点）"""
    return os.path.splitext(filename)[1][1:].lower()

def create_template_docx(template_path="template.docx"):
    """创建一个模板DOCX文件"""
    from docx import Document
    from docx.shared import Pt, RGBColor
    from docx.oxml.ns import qn
    from docx.enum.style import WD_STYLE_TYPE
    from docx.enum.table import WD_TABLE_ALIGNMENT
    
    doc = Document()
    
    # 定义要设置的所有样式名称
    style_names = [
        "Normal", "Heading 1", "Heading 2", "Heading 3", "Heading 4",
        "List Bullet", "List Number", "Table Grid", "Caption",
        "Body Text", "Quote", "Intense Quote", "Subtitle",
        "List Paragraph", "TOC 1", "TOC 2", "TOC 3"
    ]
    
    # 为所有样式设置默认格式
    for style_name in style_names:
        if style_name not in doc.styles:
            if "TOC" in style_name:
                style = doc.styles.add_style(style_name, WD_STYLE_TYPE.PARAGRAPH)
            elif style_name == "Table Grid":
                style = doc.styles.add_style(style_name, WD_STYLE_TYPE.TABLE)
            else:
                try:
                    style = doc.styles[style_name]
                except:
                    style = doc.styles.add_style(style_name, WD_STYLE_TYPE.PARAGRAPH)
        else:
            style = doc.styles[style_name]
            
        # 设置基本字体属性
        if hasattr(style, 'font'):
            style.font.name = "Times New Roman"
            style._element.rPr.rFonts.set(qn("w:eastAsia"), "宋体")
            style.font.color.rgb = RGBColor(0, 0, 0)
            
            # 根据样式类型设置不同的字号和粗细
            if "Heading 1" in style_name:
                style.font.size = Pt(16)
                style.font.bold = True
            elif "Heading 2" in style_name:
                style.font.size = Pt(14)
                style.font.bold = True
            elif "Heading 3" in style_name:
                style.font.size = Pt(13)
                style.font.bold = True
            elif "Heading 4" in style_name:
                style.font.size = Pt(12)
                style.font.bold = True
            else:
                style.font.size = Pt(12)
                style.font.bold = False
    
    # 添加示例表格以设置默认网格线
    table = doc.add_table(rows=1, cols=1)
    table.style = "Table Grid"
    table.alignment = WD_TABLE_ALIGNMENT.CENTER
    
    # 设置表格样式
    for row in table.rows:
        for cell in row.cells:
            cell.paragraphs[0].add_run("Sample")
            for border in ["top", "left", "bottom", "right"]:
                tc = cell._tc
                tc.get_or_add_tcPr()
                tcBorders = tc.get_or_add_tcPr().get_or_add_tcBorders()
                border_element = getattr(tcBorders, f"get_or_add_{border}")()
                border_element.val = "single"
                border_element.sz = 4
                border_element.color = "000000"
    
    doc.save(template_path)
    print(f"模板已保存到 {template_path}")
    return template_path

def extract_from_pdf(input_file, output_md):
    """从PDF提取文本到Markdown格式"""
    try:
        # 使用MarkItDown从PDF提取Markdown
        result = markitdown.convert(input_file)
        with open(output_md, "w", encoding="utf-8") as f:
            f.write(result)
        return output_md
    except Exception as e:
        print(f"从PDF提取文本时出错: {e}")
        return None

def convert_file(input_file, output_file, output_format, template_path=None):
    """转换单个文件"""
    input_ext = get_file_extension(input_file)
    
    # 处理PDF输入（需要特殊处理）
    if input_ext.lower() == PDF_FORMAT:
        # PDF需要先转换为Markdown
        temp_md = f"{os.path.splitext(output_file)[0]}_temp.md"
        md_file = extract_from_pdf(input_file, temp_md)
        if not md_file:
            return False, f"无法从PDF提取文本: {input_file}"
        
        # 然后从Markdown转换为目标格式
        try:
            extra_args = []
            if output_format.lower() == "docx" and template_path:
                extra_args = [
                    "--reference-doc", template_path,
                    "--toc",
                    "-f", "markdown+pipe_tables+fenced_code_blocks+east_asian_line_breaks",
                    "--wrap=none",
                    "--toc-depth=3",
                    "--columns=10000"
                ]
            
            pypandoc.convert_file(
                md_file,
                output_format,
                outputfile=output_file,
                extra_args=extra_args
            )
            
            # 删除临时Markdown文件
            if os.path.exists(temp_md):
                os.remove(temp_md)
                
            return True, f"成功转换: {input_file} -> {output_file}"
        except Exception as e:
            return False, f"转换文件时出错: {e}"
    
    # 处理非PDF输入
    else:
        try:
            extra_args = []
            if output_format.lower() == "docx" and template_path:
                extra_args = [
                    "--reference-doc", template_path,
                    "--toc",
                    "-f", f"{input_ext}+pipe_tables+fenced_code_blocks+east_asian_line_breaks",
                    "--wrap=none",
                    "--toc-depth=3",
                    "--columns=10000"
                ]
            
            pypandoc.convert_file(
                input_file,
                output_format,
                outputfile=output_file,
                extra_args=extra_args
            )
            return True, f"成功转换: {input_file} -> {output_file}"
        except Exception as e:
            return False, f"转换文件时出错: {e}"

def convert_folder(input_folder, output_folder, input_format=None, output_format="md", use_template=True):
    """转换文件夹中的所有文件"""
    results = []
    
    # 确保输出文件夹存在
    os.makedirs(output_folder, exist_ok=True)
    
    # 创建模板（如果需要）
    template_path = None
    if use_template and output_format.lower() == "docx":
        template_path = os.path.join(output_folder, "template.docx")
        create_template_docx(template_path)
    
    # 获取要转换的文件列表
    if input_format:
        input_files = glob.glob(os.path.join(input_folder, f"*.{input_format}"))
    else:
        # 如果未指定输入格式，则获取所有支持的格式
        input_files = []
        for fmt in SUPPORTED_INPUT_FORMATS + [PDF_FORMAT]:
            input_files.extend(glob.glob(os.path.join(input_folder, f"*.{fmt}")))
    
    # 转换每个文件
    for input_file in input_files:
        file_name = os.path.basename(input_file)
        file_base = os.path.splitext(file_name)[0]
        output_file = os.path.join(output_folder, f"{file_base}.{output_format}")
        
        success, message = convert_file(input_file, output_file, output_format, template_path)
        results.append({
            "input_file": input_file,
            "output_file": output_file,
            "success": success,
            "message": message
        })
    
    return results

@router.post("/convert-file/")
async def convert_single_file(
    file: UploadFile = File(...),
    output_format: str = Form(...),
    use_template: bool = Form(False)
):
    """转换单个上传的文件"""
    # 检查输出格式是否支持
    if output_format.lower() not in SUPPORTED_OUTPUT_FORMATS:
        return JSONResponse(
            status_code=400,
            content={"error": f"不支持的输出格式: {output_format}"}
        )
    
    # 创建临时目录存储上传的文件
    temp_dir = "temp_uploads"
    os.makedirs(temp_dir, exist_ok=True)
    
    # 保存上传的文件
    input_file_path = os.path.join(temp_dir, file.filename)
    with open(input_file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    
    # 设置输出文件路径
    output_dir = "output_files"
    os.makedirs(output_dir, exist_ok=True)
    file_base = os.path.splitext(file.filename)[0]
    output_file_path = os.path.join(output_dir, f"{file_base}.{output_format}")
    
    # 创建模板（如果需要）
    template_path = None
    if use_template and output_format.lower() == "docx":
        template_path = os.path.join(output_dir, "template.docx")
        create_template_docx(template_path)
    
    # 转换文件
    success, message = convert_file(input_file_path, output_file_path, output_format, template_path)
    
    # 返回结果
    return {
        "success": success,
        "message": message,
        "output_file": output_file_path if success else None
    }

@router.post("/convert-folder/")
async def convert_folder_endpoint(
    input_folder: str = Form(...),
    output_folder: str = Form(...),
    input_format: Optional[str] = Form(None),
    output_format: str = Form(...),
    use_template: bool = Form(False)
):
    """转换文件夹中的所有文件"""
    # 检查输入文件夹是否存在
    if not os.path.exists(input_folder) or not os.path.isdir(input_folder):
        return JSONResponse(
            status_code=400,
            content={"error": f"输入文件夹不存在: {input_folder}"}
        )
    
    # 检查输出格式是否支持
    if output_format.lower() not in SUPPORTED_OUTPUT_FORMATS:
        return JSONResponse(
            status_code=400,
            content={"error": f"不支持的输出格式: {output_format}"}
        )
    
    # 如果指定了输入格式，检查是否支持
    if input_format and input_format.lower() not in SUPPORTED_INPUT_FORMATS + [PDF_FORMAT]:
        return JSONResponse(
            status_code=400,
            content={"error": f"不支持的输入格式: {input_format}"}
        )
    
    # 转换文件夹
    results = convert_folder(input_folder, output_folder, input_format, output_format, use_template)
    
    # 返回结果
    return {
        "success": True,
        "results": results,
        "total_files": len(results),
        "successful_conversions": sum(1 for r in results if r["success"])
    }

@router.get("/supported-formats/")
async def get_supported_formats():
    """获取支持的文件格式"""
    return {
        "input_formats": SUPPORTED_INPUT_FORMATS + [PDF_FORMAT],
        "output_formats": SUPPORTED_OUTPUT_FORMATS
    }

# 命令行接口
def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="文档格式转换工具")
    parser.add_argument("--input", "-i", required=True, help="输入文件或文件夹路径")
    parser.add_argument("--output", "-o", required=True, help="输出文件或文件夹路径")
    parser.add_argument("--format", "-f", required=True, help="输出格式")
    parser.add_argument("--input-format", help="输入格式（仅在转换文件夹时使用）")
    parser.add_argument("--template", action="store_true", help="使用模板（用于DOCX输出）")
    
    args = parser.parse_args()
    
    # 检查输入路径是文件还是文件夹
    if os.path.isfile(args.input):
        # 单个文件转换
        output_dir = os.path.dirname(args.output)
        os.makedirs(output_dir, exist_ok=True)
        
        # 创建模板（如果需要）
        template_path = None
        if args.template and args.format.lower() == "docx":
            template_path = os.path.join(output_dir, "template.docx")
            create_template_docx(template_path)
        
        success, message = convert_file(args.input, args.output, args.format, template_path)
        print(message)
    
    elif os.path.isdir(args.input):
        # 文件夹转换
        results = convert_folder(args.input, args.output, args.input_format, args.format, args.template)
        
        # 打印结果
        print(f"总文件数: {len(results)}")
        print(f"成功转换: {sum(1 for r in results if r['success'])}")
        print(f"失败转换: {sum(1 for r in results if not r['success'])}")
        
        # 打印失败的转换
        failed = [r for r in results if not r["success"]]
        if failed:
            print("\n失败的转换:")
            for f in failed:
                print(f"  - {f['input_file']}: {f['message']}")
    
    else:
        print(f"错误: 输入路径不存在 {args.input}")

if __name__ == "__main__":
    main() 